package top.callback.inf.beautyleaf.snowflake;

import com.alibaba.fastjson.JSON;
import lombok.Data;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.FileUtils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.RetryUntilElapsed;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import top.callback.inf.beautyleaf.common.thread.MonitoredRunnable;
import top.callback.inf.beautyleaf.common.thread.NamedThreadFactory;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

/**
 * BeautyLeaf 实例管理器。一个实例就是一个 BeautyLeaf 进程，实例管理器负责管理实例在 Zookeeper 上的注册并拿到 workerID。
 * 所有的实例都会注册到 ZK_ROOT_PATH 路径下，如果这个路径不存在就表示集群第一次启动，之前没有任何实例注册过。如果 ZK_ROOT_PATH
 * 存在，那么在实例启动时就需要先去读取下面所有的 key 来判断当前实例是否曾经注册过。如果当前实例注册过，则取到 workerID 启动，
 * 否则执行一次注册。当 Zookeeper 集群无法提供服务时，会在本地查找 WORKER_ID_LOCAL_CACHE_FILE 这份缓存文件读取里面
 * 的 workerID，如果这份文件不存在，则实例无法启动。
 *
 * @author Pengyu Gan
 * CreateDate 2021/3/24
 */
@Slf4j
public class InstanceRegisterManager {

    /**
     * 本机 IP 地址
     */
    private String localIP;

    /**
     * 当前 BeautyLeaf 实例的名称
     */
    private String instanceName;

    /**
     * 当前 BeautyLeaf 实例的端口，这个端口并不是一个网络意义的端口，只是为了一台机器启动多个实例时在 ZK 里区分不同的 PATH
     */
    private int instancePort;

    /**
     * 当前 BeautyLeaf 实例在 Zookeeper 上注册的全路径
     */
    private String instancePath;

    /**
     * 当前实例地址，localIP:nodePort
     */
    private String instanceAddress;

    /**
     * Zookeeper 连接地址
     */
    private String zookeeperAddress;

    /**
     * Zookeeper 连接超时时间
     */
    private int zookeeperConnectionTimeout;

    /**
     * Zookeeper 会话超时时间
     */
    private int zookeeperSessionTimeout;

    /**
     * Snowflake 节点 ID
     */
    @Getter
    private int workerID;

    /**
     * 上一次上报（把实例信息保存到 Zookeeper 节点上）当前实例时间的时间点
     */
    private long lastUpdateTime;

    /**
     * Zookeeper 注册节点路径前缀
     */
    private final String ZK_PATH_PREFIX;

    /**
     * 保存持久数据的根节点
     */
    private final String ZK_ROOT_PATH;

    /**
     * workerID 本地缓存文件
     */
    private final String WORKER_ID_LOCAL_CACHE_FILE;

    public InstanceRegisterManager(String localIP, String instanceName, int instancePort, String zookeeperAddress, int zookeeperConnectionTimeout, int zookeeperSessionTimeout) {
        this.localIP = localIP;
        this.instanceName = instanceName;
        this.instancePort = instancePort;
        this.instanceAddress = localIP + ":" + instancePort;
        this.zookeeperAddress = zookeeperAddress;
        this.zookeeperConnectionTimeout = zookeeperConnectionTimeout;
        this.zookeeperSessionTimeout = zookeeperSessionTimeout;
        ZK_PATH_PREFIX = "/snowflake/" + instanceName;
        ZK_ROOT_PATH = ZK_PATH_PREFIX + "/forever";
        WORKER_ID_LOCAL_CACHE_FILE = System.getProperty("java.io.tmpdir") + File.separator + instanceName + "/leafconf/%s/workerID.properties";
    }

//    public InstanceRegisterManager(String localIP, int instancePort, String zookeeperAddress) {
//        this.localIP = localIP;
//        this.instancePort = instancePort;
//        this.instanceAddress = localIP + ":" + instancePort;
//        this.zookeeperAddress = zookeeperAddress;
//    }

    public boolean init() {
        try {
            // 实例化 Curator
            CuratorFramework curator = createCurator();
            curator.start();
            // 检查根节点是否存在
            Stat stat = curator.checkExists().forPath(ZK_ROOT_PATH);
            // 根节点不存在表示第一次启动，为当前实例创建永久节点 /snowflake/Chengdu/forever/本机IP:节点端口-序列号(0000000000)
            if (null == stat) {
                // 在 Zookeeper 上创建当前 BeautyLeaf 实例对应的节点并拿到节点全路径
                instancePath = createPersistentNodeForInstance(curator);
                // 缓存 workerID 到本地文件
                updateOrCreateWorkerIDLocalCacheFile(workerID);
                // 定时将当前实例的本机时间保存到 Zookeeper
                scheduleUpdateInstanceTime(curator, instancePath);
                return true;
            } else { // 根节点存在，表示肯定有某个实例启动过，由于不确定当前实例是否启动过，做以下操作进行检查
                // 拿到根路径下注册的所有 BeautyLeaf 实例的 key，格式为 instanceAddress-序列号，如 172.16.0.110:8888-000000001
                List<String> keysInRootPath = curator.getChildren().forPath(ZK_ROOT_PATH);
                // instanceAddress -> workerID(Zookeeper 为节点生成的序列号)
                Map<String, Integer> instanceAddressToWorkerID = new HashMap<>();
                // instanceAddress -> 实例的 key
                Map<String, String> instanceAddressToKey = new HashMap<>();
                for (String key : keysInRootPath) {
                    String[] keyPair = key.split("-");
                    String instanceAddress = keyPair[0];
                    String nodeSequence = keyPair[1];
                    instanceAddressToWorkerID.put(instanceAddress, Integer.valueOf(nodeSequence));
                    instanceAddressToKey.put(instanceAddress, key);
                }
                // 拿到当前实例的 workerID
                Integer workerID = instanceAddressToWorkerID.get(instanceAddress);
                // 如果 workerID 存在，表示曾经注册过
                if (null != workerID) {
                    // 拿到当前实例的注册路径
                    instancePath = ZK_ROOT_PATH + "/" + instanceAddressToKey.get(instanceAddress);
                    // 保存 workerID 到当前实例
                    this.workerID = workerID;
                    boolean isClockBack = checkClockBack(curator, instancePath);
                    if (isClockBack) {
                        throw new RuntimeException("当前实例所在机器时间小于上一次注册时间，可能发生了时钟回拨。");
                    }
                    // 定时将当前实例的本机时间保存到 Zookeeper
                    scheduleUpdateInstanceTime(curator, instancePath);
                    // 缓存 workerID 到本地文件
                    updateOrCreateWorkerIDLocalCacheFile(this.workerID);
                    log.info("已找到当前实例曾经的注册信息，当前实例启动成功。实例 IP: {}, 实例端口: {}, workerID: {}", localIP, instancePort, this.workerID);
                } else { // workerID 不存在，当前实例没有在 Zookeeper 注册过
                    // 在 Zookeeper 上创建当前 BeautyLeaf 实例对应的节点
                    instancePath = createPersistentNodeForInstance(curator);
                    // 从 instancePath 上拿到 Zookeeper 生成的节点编号作为 workerID
                    this.workerID = Integer.parseInt(instancePath.split("-")[1]);
                    // 定时将当前实例的本机时间保存到 Zookeeper
                    scheduleUpdateInstanceTime(curator, instancePath);
                    // 缓存 workerID 到本地文件
                    updateOrCreateWorkerIDLocalCacheFile(this.workerID);
                    log.info("新实例注册并启动成功。实例 IP: {}, 实例端口: {}, workerID: {}", localIP, instancePort, this.workerID);
                }
            }
        } catch (Exception e) {
            log.error("当前实例启动失败，尝试查找本地 workerID 缓存文件进行启动。", e);
            try {
                // Zookeeper 集群出现问题无法连接，尝试从本地缓存文件中找到 workerID 启动
                Properties properties = new Properties();
                properties.load(new FileInputStream(String.format(WORKER_ID_LOCAL_CACHE_FILE, instancePort)));
                this.workerID = Integer.parseInt(properties.getProperty("workerID"));
                log.info("当前实例使用本地缓存的 workerID {} 启动。", this.workerID);
            } catch (Exception ex) {
                // 本地不存在缓存文件，实例无法启动
                log.error("当前实例无法找到本地缓存的 workerID，彻底无法启动。", ex);
                return false;
            }
        }
        return true;
    }

    /**
     * 检查当前实例所在机器是否发生了时钟回拨
     * @param curator Zookeeper 操作接口
     * @param instancePath 当前实例曾经在 Zookeeper 上的注册地址
     * @return true 时间回退 false 时间正常
     * @throws Exception 获取 Zookeeper 上数据异常时抛出
     */
    private boolean checkClockBack(CuratorFramework curator, String instancePath) throws Exception {
        // 从 Zookeeper 拿回注册的实例数据
        byte[] data = curator.getData().forPath(instancePath);
        Instance instance = JSON.parseObject(data, Instance.class);
        long now = System.currentTimeMillis();
        // 时钟回拨打印两个时间方便排查问题
        if (instance.getTimestamp() > now) {
            log.error("当前实例所在机器时间 {} 小于上一次注册时间 {}，可能发生了时钟回拨。", instance.getTimestamp(), now);
        }
        return instance.getTimestamp() > now;
    }

    /**
     * 定时(3 秒)将当前实例的本机时间保存到 Zookeeper
     * @param curator Zookeeper 操作接口
     * @param instancePath 数据保存路径
     */
    private void scheduleUpdateInstanceTime(CuratorFramework curator, String instancePath) {
        Executors.newSingleThreadScheduledExecutor(new NamedThreadFactory("BeautyLeaf-Instance-Time-Update", true))
                .scheduleWithFixedDelay(new MonitoredRunnable(
                        () -> updateInstanceTime(curator, instancePath),
                        t -> log.error("定时任务 更新当前 BeautyLeaf 实例本机时间到 Zookeeper 异常", t)),
                        1L, 3L, TimeUnit.SECONDS);
    }

    /**
     * 更新当前 BeautyLeaf 实例本机时间到 Zookeeper
     * @param curator Zookeeper 操作接口
     * @param instancePath 数据保存路径
     */
    private void updateInstanceTime(CuratorFramework curator, String instancePath) {
        try {
            // 当前时间小于上一次更新时间，表示出现了时钟回退，放弃本次上报
            if (System.currentTimeMillis() < lastUpdateTime) {
                return;
            }
            // 构建实例数据上报
            String data = buildInstanceInfo();
            curator.setData().forPath(instancePath, data.getBytes(StandardCharsets.UTF_8));
        } catch (Exception e) {
            log.error("更新当前 BeautyLeaf 实例本机时间到 Zookeeper 异常", e);
        }
    }

    /**
     * 在 Zookeeper 上创建当前 BeautyLeaf 实例对应的节点
     * @param curator zookeeper 操作接口
     * @return 当前 BeautyLeaf 实例对应节点的全路径
     * @throws Exception 创建失败时抛出此异常
     */
    private String createPersistentNodeForInstance(CuratorFramework curator) throws Exception {
        try {
            // 实例对应节点的全路径，Zookeeper 会默认在后面加上序列号，这里用 - 号分割，方便之后用于 workerID
            String path = ZK_ROOT_PATH + "/" + instanceAddress + "-";
            // 构建实例数据
            String data = buildInstanceInfo();
            return curator.create()
                    .creatingParentsIfNeeded()
                    .withMode(CreateMode.PERSISTENT_SEQUENTIAL)
                    .forPath(path, data.getBytes(StandardCharsets.UTF_8));
        } catch (Exception e) {
            log.error("创建当前 BeautyLeaf 节点在 Zookeeper 上对应节点异常", e);
            throw e;
        }
    }

    /**
     * 构建当前 BeautyLeaf 实例的信息
     * @return JSON 格式的当前 BeautyLeaf 实例信息
     */
    private String buildInstanceInfo() {
        Instance instance = new Instance(localIP, String.valueOf(instancePort), System.currentTimeMillis());
        return JSON.toJSONString(instance);
    }

    /**
     * 更新或创建 workerID 本地缓存文件
     * @param workerID 节点 ID
     */
    private void updateOrCreateWorkerIDLocalCacheFile(int workerID) {
        File workerIDLocalCacheFile = new File(String.format(WORKER_ID_LOCAL_CACHE_FILE, instancePort));
        boolean exists = workerIDLocalCacheFile.exists();
        // 文件存在则更新
        if (exists) {
            try {
                // 覆盖写入 workerID
                FileUtils.writeStringToFile(workerIDLocalCacheFile, "workerID=" + workerID, StandardCharsets.UTF_8, false);
                log.info("更新 workerID {} 到本地缓存文件成功。", workerID);
            } catch (IOException e) {
                log.error("更新 workerID {} 到本地文件异常", workerID, e);
            }
        } else { // 文件不存在则先尝试新建
            try {
                // 建立父目录
                boolean createParentFileSuccess = workerIDLocalCacheFile.getParentFile().mkdirs();
                if (createParentFileSuccess) {
                    // 文件创建成功，写入 workerID
                    if (workerIDLocalCacheFile.createNewFile()) {
                        FileUtils.writeStringToFile(workerIDLocalCacheFile, "workerID=" + workerID, StandardCharsets.UTF_8, false);
                        log.info("workerID {} 写入本地缓存文件成功。", workerID);
                    }
                } else {
                    log.error("创建 workerID {} 本地缓存文件的父目录失败", workerID);
                }
            } catch (IOException e) {
                log.error("创建 workerID {} 本地缓存文件失败", workerID, e);
            }
        }
    }

    /**
     * 创建 CuratorFramework 实例以操作 Zookeeper
     * @return CuratorFramework 实例
     */
    private CuratorFramework createCurator() {
        return CuratorFrameworkFactory.builder()
                .connectString(zookeeperAddress)
                .retryPolicy(new RetryUntilElapsed(1000, 4))
                .connectionTimeoutMs(zookeeperConnectionTimeout)
                .sessionTimeoutMs(zookeeperSessionTimeout)
                .build();
    }

    /**
     * BeautyLeaf 实例信息，一个实例就是一个 BeautyLeaf 进程。
     */
    @Data
    static class Instance {

        /**
         * 本机 IP
         */
        private String localIP;

        /**
         * 当前 BeautyLeaf 节点的端口
         */
        private String instancePort;

        /**
         * 当前 BeautyLeaf 节点的本机时间
         */
        private long timestamp;

        public Instance() {
        }

        public Instance(String localIP, String instancePort, long timestamp) {
            this.localIP = localIP;
            this.instancePort = instancePort;
            this.timestamp = timestamp;
        }
    }

}
